#!/usr/bin/env python
from __future__ import print_function
import os
import re
import numpy as np

try:
    from astropy.io import fits as pf
except ImportError:
    import pyfits as pf

class SimpleEBoundObj(object):

    def __init__(self, infile):
        if not isinstance(infile, str):
            raise IOError('infile should be a string!')
        self.infile = infile
        hdulist = pf.open(self.infile, memmap=True)
        try:
            EBoundHDU = hdulist['EBOUNDS']
            self.Emin = np.array(EBoundHDU.data.field(1) * 1.e-3).flatten() # [MeV]
            self.Emax = np.array(EBoundHDU.data.field(2) * 1.e-3).flatten() # [MeV]
            self.Energy = np.unique(np.r_[self.Emin, self.Emax]) # [MeV]
            self.ebType = 'ccube'
        except KeyError:
            EBoundHDU = hdulist['ENERGIES']
            self.Energy = EBoundHDU.data.field(0) # [MeV]
            self.Emin = np.array(self.Energy[:-1]).flatten() # [MeV]
            self.Emax = np.array(self.Energy[1:]).flatten()  # [MeV]
            self.ebType = 'expcube'
        hdulist.close()
        del EBoundHDU.data

    @property
    def nEBound(self):
        return len(self.Emin) # No. of Ebins

    @property
    def nEnergy(self):
        return len(self.Energy) # No. of energy boundary

formatter_defaults = [
    'results_pyLike_band{index:02d}.dat', # new pyLike
    'results_garLike_band{index:02d}.dat', # gardian
    'results_pyLike_band{index:d}.dat',   # old pyLike
    'results_pyLike_band{index:02d}_wid02.dat', # binwd 2
    'results_pyLike_band{index:02d}_wid03.dat',
    'results_pyLike_band{index:02d}_wid04.dat',
    'results_pyLike_band{index:02d}_wid05.dat',
    'results_pyLike.dat',
    'results_garLike_band{index:02d}_wid02.dat', # binwd 2
    'results_garLike_band{index:02d}_wid03.dat',
    'results_garLike_band{index:02d}_wid04.dat',
    'results_garLike_band{index:02d}_wid05.dat',
    ]

def collect_data(centerSrc, ccube=None, formatter=None, flux_factor=1.):
    eb = None
    com_data = {'emins': [], 'emaxs': [], 'logLike': [],
                'fitqual': [], 'retcode': [], 'edm': []}
    target_srcs = set([])
    src_data = {}
    for iband in range(200):
        # load file
        if formatter is None or not os.path.isfile(formatter.format(index=iband)):
            for fmter in formatter_defaults:
                if os.path.isfile(fmter.format(index=iband)):
                    formatter = fmter
                    break
            if formatter is None:
                continue

        datfile = formatter.format(index=iband)
        if iband>0 and datfile==formatter:
            break

        try:
            datafh = open(datfile, 'r')
        except IOError: # as err:
            continue
        else:
            info = eval(datafh.read())
            datafh.close()

        # collect common data
        if eb is None:
            if 'zERange' in info:
                _emin, _emax = info['zERange']
            elif '+ERange' in info:
                _emin, _emax = info['+ERange']
            elif ccube is None:
                raise RuntimeError('No energy info in the dat file? ccube is needed!')
            else:
                eb = SimpleEBoundObj(ccube)
                _emin, _emax = eb.Emin[iband], eb.Emax[iband]
        else:
            _emin, _emax = eb.Emin[iband], eb.Emax[iband]

        com_data['emins'].append(_emin)
        com_data['emaxs'].append(_emax)

        if 'zFitQuality' in info:
            com_data['fitqual'].append(info['zFitQuality'])
        else:
            com_data['fitqual'].append(-1)

        if 'zReturnCode' in info:
            com_data['retcode'].append(info['zReturnCode'])
        else:
            com_data['retcode'].append(-1)

        if 'zEDM' in info:
            com_data['edm'].append(info['zEDM'])
        elif '+EDM' in info:
            com_data['edm'].append(info['+EDM'])
        else:
            com_data['edm'].append(-1)

        if 'logLikelihood' in info:
            com_data['logLike'].append(info['logLikelihood'])
        else:
            com_data['logLike'].append(0.)

        # collect target source names
        if not target_srcs:
            if isinstance(centerSrc, str):
                snames0 = [centerSrc]
            elif isinstance(centerSrc, (list, tuple, set)):
                snames0 = centerSrc
            else:
                raise IOError('Invalid centerSrc')

            for sname0 in snames0:
                if not sname0.strip() or sname0.startswith('#'):
                    continue
                if sname0 in info:
                    target_srcs.add(sname0)
                else:
                    for sname in info:
                        if re.search(sname0, sname):
                            target_srcs.add(sname)

            if target_srcs:
                print('iBin\tflux/fluxE\tsqrt(npred)\tflux/(fluxE*sqrt(npred))')
            else:
                raise RuntimeError('No matched source found!')

        # collect data for each data
        for target in target_srcs:
            sinfo = info[target]
            if isinstance(sinfo, list):
                sinfo = sinfo[1]
            if not isinstance(sinfo, dict):
                continue

            if target in src_data:
                sdata = src_data[target]
            else:
                sdata = {'flux': [], 'fluxerr': [], 'fluxul': [], 'ts': [], 'npred': []}

            try:
                myts = float(sinfo['TS value'])
            except:
                myts = float(sinfo['+TSValue'])
            sdata['ts'].append(myts)

            try:
                npred = float(sinfo['NPred'])
            except KeyError:
                npred = float(sinfo['+NPred'])
            sdata['npred'].append(npred)

            if '+Flux' in sinfo:
                if '+/-' in sinfo['+Flux']:
                    myflux, myfluxerr = [float(v) for v in sinfo['+Flux'].split('+/-')]
                else:
                    myflux, myfluxerr = float(sinfo['+Flux']), 0.
            else:
                if '+/-' in sinfo['Flux']:
                    myflux, myfluxerr = [float(v) for v in sinfo['Flux'].split('+/-')]
                else:
                    myflux, myfluxerr = float(sinfo['Flux']), 0.

            if isinstance(flux_factor, (list, np.ndarray)):
                flux_factor_ = flux_factor[iband]
            else:
                flux_factor_ = flux_factor
            myflux *= flux_factor_
            myfluxerr *= flux_factor_

            if 'Upper Limit' in sinfo:
                myfluxul = float(sinfo['Upper Limit'])*flux_factor_
            elif '+UpperLimit' in sinfo:
                myfluxul = float(sinfo['+UpperLimit'])*flux_factor_
            else:
                myfluxul = -1.

            sdata['flux'].append(myflux)
            sdata['fluxerr'].append(myfluxerr)
            sdata['fluxul'].append(myfluxul)

            # print out some check values
            if myfluxul > 0.:
                print('{iband:02d}[{sname:s}]: {ts:.3f} [UL]'.format(iband=iband, sname=target, ts=myts))
            else:
                if npred < 0. or myfluxerr == 0.:
                    print('{iband:02d}[{sname:s}]: skip'.format(iband=iband, sname=target))
                else:
                    f_frac = myflux/myfluxerr
                    n_frac = np.sqrt(npred)
                    print('{iband:02d}[{sname:s}]: {f_frac:.6f}\t{n_frac:.6f}\t{fn_frac_rat:.6f}'.format(
                            iband=iband, sname=target, f_frac=f_frac, n_frac=n_frac, fn_frac_rat=f_frac/n_frac
                        )
                    )

            src_data[target] = sdata

    com_data['srcdata'] = src_data
    return com_data

def build_plotdata(indata, ULTSLimit=None):
    try:
        tslim = float(ULTSLimit)
    except:
        tslim = -float('inf')

    src_data = indata['srcdata']

    plotdata = {}
    for sname, sdata in src_data.items():
        flux, fluxerr, ulflags = [], [], []
        for fl, fe, ul, ts in zip(sdata['flux'], sdata['fluxerr'], sdata['fluxul'], sdata['ts']):
            if ul > 0. and ts < tslim:
                flux.append(ul)
                fluxerr.append(0.)
                ulflags.append(-1.)
            else:
                flux.append(fl)
                fluxerr.append(fe)
                ulflags.append(0.)

        outdata = np.c_[
            indata['emins'], indata['emaxs'],
            np.atleast_1d(flux), np.atleast_1d(fluxerr), sdata['npred'], sdata['ts'],
            indata['fitqual'], indata['retcode'], indata['edm'],
            np.atleast_1d(ulflags), indata['logLike']
        ]

        plotdata[sname] = outdata

    return plotdata

def store_data(indata, ULTSLimit=None, output='myseddata_{sname:s}.dat'):
    plotdata = build_plotdata(indata, ULTSLimit=ULTSLimit)
    for sname, dat in plotdata.items():
        sname_ = '_'.join(sname.split())
        if '%s' in output:
            outfile = output%sname_
        else:
            outfile = output.format(sname=sname_)

        with open(outfile, 'w') as fh:
            fh.write('# Emin[MeV]\tEmax[MeV]\tflux(_UL)[ph cm^-2 s^-1]\tfluxErr'+\
                    '[ph cm^-2 s^-1]\tnpred\tts\tfitQual\tretCode\tEDM\tLIMIT_FLAG\tlogLike\n')
            np.savetxt(fh, dat, delimiter='\t', fmt=['% .6e']*4+['% .4e']*2+['% 4d']*2 + ['% .1e', '% d', '% e'])

    return plotdata

def plotSED_data(plotdata, outfile='plot_sed.pdf',
    fig=None, alpha=2, plotdict=None, xlim=None, ylim=None, y2lim=None, remove_xerr=False, xoffset_fac=1.):
    from matplotlib import pyplot as plt

    if fig is None:
        fig = plt.gcf()

    if fig.axes:
        if len(fig.axes)==1:
            ax1 = fig.axes[0]
            ax2 = None
        else:
            ax1, ax2 = fig.axes[-2:]
    else:
        ax1 = fig.add_subplot(111)
        ax2 = ax1.twinx()

    if not ax1.get_xlabel():
        ax1.set_xlabel(r'$E_{\gamma}$ (${\rm GeV}$)')
        ax1.set_xscale('log')
    if not ax1.get_ylabel():
        ax1.set_ylabel(
            r'$E_{\gamma}^{%s}\ {\rm d}N/{\rm d}E_{\gamma}$ (${\rm GeV^{%s}\ ph\ cm^{-2}\ s^{-1}\ GeV^{-1}}$)'
            %(alpha, alpha)
            )
        ax1.set_yscale('log')
    if ax2 and not ax2.get_ylabel():
        ax2.set_ylabel('TS Value')

    emin, emax = plotdata[:, 0]*1e-3, plotdata[:, 1]*1e-3
    emid = np.sqrt(emin * emax) # GeV
    e2dnde = plotdata[:, 2]/(emax-emin)*emid**alpha
    e2dnde_err = plotdata[:, 3]/(emax-emin)*emid**alpha
    tsval = plotdata[:, 5]
    ulflag = plotdata[:, 9]

    ulma = ulflag < 0.

    cstr = 'C1'
    default_dict = dict(ms=3)
    default_dict_ts = dict() # for TS values
    if not plotdict is None:
        default_dict.update(plotdict)
        if 'marker' in plotdict:
            del default_dict['marker']
        if 'ls' in plotdict:
            del default_dict['ls']
        if 'hatch' in plotdict:
            default_dict_ts['hatch'] = default_dict.pop('hatch')

        if 'c' in plotdict:
            cstr = plotdict['c']
        elif 'ecolor' in plotdict:
            cstr = plotdict['ecolor']
        elif 'mfc' in plotdict:
            cstr = plotdict['mfc']
        elif 'mec' in plotdict:
            cstr = plotdict['mec']
    default_dict['c'] = cstr
    default_dict['ecolor'] = cstr
    default_dict['mec'] = cstr
    default_dict['mfc'] = cstr

    default_dict2 = default_dict.copy()
    if plotdict is None:
        default_dict2['marker'] = 'o'
        default_dict2['ls'] = 'None'
    else:
        default_dict2['marker'] = plotdict.get('marker', 'o')
        default_dict2['ls'] = plotdict.get('ls', 'None')
    if remove_xerr:
        ax1.errorbar(emid[~ulma]*xoffset_fac, e2dnde[~ulma], yerr=e2dnde_err[~ulma], **default_dict2)
    else:
        ax1.errorbar(emid[~ulma]*xoffset_fac, e2dnde[~ulma], yerr=e2dnde_err[~ulma], xerr=((emid*xoffset_fac-emin)[~ulma], (emax-emid*xoffset_fac)[~ulma]), **default_dict2)

    if ulma.all():
        mymax = -np.inf
        mymin = np.inf
    else:
        mymax = (e2dnde+e2dnde_err)[~ulma].max()
        mymin = (e2dnde-e2dnde_err)[~ulma].min()
    if ulma.any():
        mymax = max(mymax, e2dnde[ulma].max())
        mymin = min(mymin, e2dnde[ulma].min()*0.5)
    assert mymax > 0. and mymax > mymin
    if mymin < 1e-5*mymax:
        mymedian = np.median(e2dnde)
        assert mymedian > 0.
        mymin = min(mymedian**2/mymax, e2dnde.min()*0.9)

    if ulma.any():
        # determine the size of upperlimit arrow
        wd, ht = fig.get_figwidth(), fig.get_figheight()
        xbinwd = np.log(emax[0]/emin[0])/np.log(emax[-1]/emin[0])/wd
        ybinfc = 1.-1./np.exp(xbinwd*ht*np.log(mymax/mymin))
        assert 0.<ybinfc<1.

        if 'label' in default_dict:
            del default_dict['label']
        if 'capsize' in default_dict:
            del default_dict['capsize']
        ax1.errorbar(emid[ulma]*xoffset_fac, e2dnde[ulma], yerr=e2dnde[ulma]*ybinfc*2., xerr=((emid*xoffset_fac-emin)[ulma], (emax-emid*xoffset_fac)[ulma]),
                    uplims=True, ls='None', **default_dict)

    if ax2:
        ax2.bar((emax+emin)*0.5, tsval, emax-emin, color='0.6', alpha=0.3, **default_dict_ts)

    if xlim is None:
        ax1.set_xlim((emin.min()*0.8, emax.max()*1.25))
        if ax2:
            ax2.set_xlim((emin.min()*0.8, emax.max()*1.25))
    else:
        ax1.set_xlim(xlim)
        if ax2:
            ax2.set_xlim(xlim)

    if ylim is None:
        lgmin, lgmax = np.log10(mymin), np.log10(mymax)
        delta = lgmax-lgmin
        if delta <= 1.5:
            lgmin -= (1.5-delta)/2.
            lgmax += (1.5-delta)/2.
        ymin_10, ymax_10 = lgmin//1., lgmax//1.
        ymin_pf, ymax_pf = 10.**(lgmin-ymin_10), 10.**(lgmax-ymax_10)
        ymin = np.round(ymin_pf-0.5)*10.**ymin_10
        ymax = np.round(ymax_pf+0.5)*10.**ymax_10
       #print 10.**lgmin, ymin, 10.**lgmax, ymax
        ax1.set_ylim((ymin, ymax))
    else:
        ax1.set_ylim(ylim)

    if ax2:
        if y2lim is None:
            ax2.set_ylim((0, np.ceil((tsval*1.1).max())))
        else:
            ax2.set_ylim(y2lim)

    #fig.tight_layout()
    if outfile:
        fig.savefig(outfile)

    if ax2:
        return fig, ax1, ax2
    else:
        return fig, ax1


# backward support
def collect(centerSrc, ccube=None, ULTSLimit=None,
            output='myseddata_{sname:s}.dat', formatter=None,
            flux_factor=1.):
    mydata = collect_data(centerSrc, ccube=ccube, formatter=formatter, flux_factor=flux_factor)
    return store_data(mydata, ULTSLimit=ULTSLimit, output=output).keys()

def plotSED(datafile, outfile='plot_sed.pdf',
            fig=None, alpha=2, plotdict=None, xlim=None, ylim=None, y2lim=None, remove_xerr=False, xoffset_fac=1.):
    if isinstance(datafile, np.ndarray):
        data = datafile
    else:
        data = np.loadtxt(datafile)
    if data.ndim==1 and data.shape[0]==11:
        data = data.reshape(1, -1)
    assert data.ndim==2 and data.shape[0]>0

    return plotSED_data(data, outfile=outfile,
            fig=fig, alpha=alpha, plotdict=plotdict, xlim=xlim, ylim=ylim, y2lim=y2lim, remove_xerr=remove_xerr, xoffset_fac=xoffset_fac)

# command line interface
def cli():
    import argparse

    parser = argparse.ArgumentParser(description='Collect SED results')
    parser.add_argument("centerSrc", type=str, help='Center source')
    parser.add_argument("-ccube", type=str, default=None,
            help='ccube file (needed when dat file do not have energy info)')
    parser.add_argument("-sedfile", type=str, default=None,
            help='the sed dat file that used to generate sed plots')
    parser.add_argument("-ULTSLimit", type=float, default=10.,
            help='TS below ULTSLimit will read Upper Limit, default: 10')
    parser.add_argument("-sedout", type=str, default='myseddata_{sname:s}.dat',
            help='The output sed filename')
    parser.add_argument("-formatter", type=str, default=None,
            help='Formatter of the results.dat')
    parser.add_argument("-fluxfactor", type=float, default=1.,
            help='The factor multiplied on the flux, default: 1.')
    parser.add_argument("-show", action='store_true', default=False, help='Whether to show the SED')
    parser.add_argument("-plot", action='store_true', default=False, help='Whether to plot the SED')
    parser.add_argument("-plotout", type=str, default='fig_sed_{sname:s}.pdf',
            help='The filename of the plotted figure')

    args = parser.parse_args()

    ## only show the figure
    import matplotlib as mpl
    if args.sedfile is not None:
        from matplotlib import pyplot as plt
        isshow = not args.plot
        if isshow:
            outfile = None
        else:
            mpl.use('Agg')
            if '%s' in args.plotout:
                outfile = args.plotout%args.centerSrc
            else:
                outfile = args.plotout.format(sname=args.centerSrc)

        plotSED(args.sedfile, outfile)
        if isshow:
            plt.show()
        return

    ## generate the dat file and show the figure
    if not args.show:
        mpl.use('Agg')

    ctrfile = None
    if args.centerSrc.startswith('@'):
        ctrfile = args.centerSrc[1:]
        assert os.path.isfile(ctrfile)
    elif os.path.isfile(args.centerSrc):
        ctrfile = args.centerSrc
    else:
        ctrsrcs = args.centerSrc
    if ctrfile is not None:
        ctrsrcs = []
        with open(ctrfile) as fh:
            for l in fh:
                if l.startswith('#'): continue
                ctrsrcs.append(l.strip())

    snames = collect(ccube=args.ccube, centerSrc=ctrsrcs, ULTSLimit=args.ULTSLimit,
                     output=args.sedout, formatter=args.formatter, flux_factor=args.fluxfactor)

    if args.plot or args.show:
        from matplotlib import pyplot as plt
        for sname in snames:
            sname_ = '_'.join(sname.split())
            if '%s' in args.sedout:
                datafile = args.sedout%sname_
            else:
                datafile = args.sedout.format(sname=sname_)

            if args.plot:
                if '%s' in args.plotout:
                    outfile = args.plotout%sname_
                else:
                    outfile = args.plotout.format(sname=sname_)
            else:
                outfile = None
            plotSED(datafile, outfile)

            if args.plot:
                plt.clf()

        if args.show:
            plt.show()

if __name__ == '__main__':
    cli()
